from fastapi import FastAPI, File, UploadFile
from PIL import Image
import torch
from torchvision import transforms
from efficientnet_pytorch import EfficientNet
import io

app = FastAPI()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = EfficientNet.from_pretrained("efficientnet-b0", num_classes=3)
model.load_state_dict(torch.load('./result/checkpoint_epoch_7.pth', map_location=device, weights_only=True))
model.to(device)
model.eval()

id2label = {0: 'cat', 1: 'dog', 2: 'other'}

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # 读取上传文件为PIL Image
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert('RGB')
    # 预处理并推理
    input_tensor = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = model(input_tensor)
        _, predicted = torch.max(outputs, dim=1)

    label = id2label[predicted.item()]
    return {"predicted_class": label}
